from omegaconf import DictConfig

import numpy as np
from torch.utils.data import Dataset

from ecgcmr.signal.sig_augmentations.ECGAugmentations import ECGAugmentations


class ContrastiveECGDataset(Dataset):
  def __init__(
      self, 
      cfg: DictConfig, 
      mode: str,
      apply_augmentations: bool = True,
      ) -> None:
    
    if mode == 'train':
      self.data_ecg = np.load(cfg.dataset.paths.data_ecg_train, mmap_mode='c')
    elif mode == 'val':
      self.data_ecg = np.load(cfg.dataset.paths.data_ecg_val, mmap_mode='c')
    elif mode == 'test':
      self.data_ecg = np.load(cfg.dataset.paths.data_ecg_test, mmap_mode='c')
    
    self.ecg_augmentations = ECGAugmentations(cfg=cfg.augmentations.ecg,
                                              sampling_rate=cfg.dataset.sampling_rate,
                                              input_electrodes=cfg.dataset.input_electrodes,
                                              single_modality=cfg.dataset.single_modality,
                                              use_peaks_location=cfg.dataset.use_peaks_location,
                                              apply_augmentations=apply_augmentations,
                                              seed=cfg.seed)
  def __len__(self):
    return len(self.data_ecg)
  
  def __getitem__(self, index):
    return self.ecg_augmentations(self.data_ecg[index]), self.ecg_augmentations(self.data_ecg[index])